#!/usr/bin/env bash
set -euo pipefail

# Usage: scripts/run_local_dumbo.sh [N=4] [DURATION=30] [BATCH_SIZE=100]

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
cd "$ROOT_DIR"

N="${1:-4}"
DURATION="${2:-30}"
BATCH_SIZE="${3:-100}"

echo "[run_local_dumbo] root: $ROOT_DIR, N=$N, duration=${DURATION}s, batch_size=$BATCH_SIZE"

# Output data root and timestamped run dir
DATA_ROOT="$ROOT_DIR/data"
TS=$(date +%Y%m%d_%H%M%S)
RUN_DIR="$DATA_ROOT/$TS"
export RUN_DIR
mkdir -p "$RUN_DIR/logs"
mkdir -p "$RUN_DIR/results"
echo "[run_local_dumbo] data dir: $RUN_DIR"

# Build Go binary
echo "[run_local_dumbo] building Go binary..."
go mod tidy
go build -o dumbo main.go

# Generate configurations for local testing
echo "[run_local_dumbo] generating configurations..."
export N_VAL="$N"
export BATCH_SIZE_VAL="$BATCH_SIZE"
export DURATION_VAL="$DURATION"
export TS
python3 - <<'PY'
import yaml
import os

N = int(os.environ.get('N_VAL','4'))
BATCH_SIZE = int(os.environ.get('BATCH_SIZE_VAL','100'))
RUN_DIR = os.environ['RUN_DIR']

# Create config template for local testing
config_template = {
    'id_name': {i: f'node{i}' for i in range(N)},
    'id_ip': {i: f'127.0.0.1' for i in range(N)},
    'id_p2p_port': {i: 9000 + i for i in range(N)},  # Different ports for each node
    'log_level': 1,
    'max_pool': 10,
    'timeout': 1000,
    'mock_latency': 10,
    'ddos': True,
    'max_payload_size': BATCH_SIZE,  # Configurable batch size
    'tx_size': 32,
    'wait_time': 3  # Reduced wait time for local testing
}

# Write config template
template_path = os.path.join(RUN_DIR, 'config_temp.yaml')
with open(template_path, 'w') as f:
    yaml.dump(config_template, f, default_flow_style=False)

print(f'Config template written to: {template_path} with batch_size={BATCH_SIZE}')
PY

# Generate individual node configurations
echo "[run_local_dumbo] generating node configurations..."
cd config_gen
cp "$RUN_DIR/config_temp.yaml" ./config_temp.yaml
go run main.go

# Fix case sensitivity issues in generated configs
for yaml_file in *.yaml; do
  if [ -f "$yaml_file" ]; then
    # Convert TSPubKey to tspubkey and TSShare to tsshare
    sed -i 's/TSPubKey:/tspubkey:/g; s/TSShare:/tsshare:/g' "$yaml_file"
  fi
done

cp *.yaml "$RUN_DIR/"
cd ..

echo "[run_local_dumbo] configurations generated"

# Cleanup old processes
pkill -f "dumbo -config" || true
sleep 1

# Start Dumbo nodes
echo "[run_local_dumbo] starting $N Dumbo node(s)..."
for (( i=0; i< N; i++ )); do
  config_file="$RUN_DIR/node$i.yaml"
  # Change to the directory containing the config file and run dumbo
  (cd "$RUN_DIR" && "$ROOT_DIR/dumbo" -config "node$i" > "$RUN_DIR/logs/node-$i.log" 2>&1) &
  echo "[run_local_dumbo] started node $i with config: $config_file (PID: $!)"
done

echo "[run_local_dumbo] running for ${DURATION}s..."
sleep "$DURATION"

echo "[run_local_dumbo] stopping nodes..."
pkill -f "dumbo -config" || true
sleep 2

echo "[run_local_dumbo] logs:"
ls -la "$RUN_DIR/logs" || true

# Parse logs for Dumbo-specific statistics
echo "[run_local_dumbo] parsing logs for statistics..."
python3 - <<'PY'
import glob, json, os
from pathlib import Path

RUN_DIR = os.environ['RUN_DIR']
logs = sorted(glob.glob(os.path.join(RUN_DIR,'logs/*.log')))
if not logs:
    print('no logs found, skip stats')
else:
    def parse_dumbo_log(path: str) -> dict:
        result = {
            'file': Path(path).name, 
            'lines': 0, 
            'rbc_outputs': 0, 
            'aba_decisions': 0, 
            'mvba_outputs': 0, 
            'acs_completions': 0,
            'consensus_rounds': 0,
            'errors': 0
        }
        try:
            with open(path, 'r', errors='ignore') as f:
                for line in f:
                    result['lines'] += 1
                    if 'RBC delivered' in line:
                        result['rbc_outputs'] += 1
                    if 'ABA decided' in line:
                        result['aba_decisions'] += 1
                    if 'MVBA delivered' in line:
                        result['mvba_outputs'] += 1
                    if 'ACS completed' in line:
                        result['acs_completions'] += 1
                    if 'Consensus round completed' in line:
                        result['consensus_rounds'] += 1
                    if 'error' in line.lower() or 'panic' in line.lower():
                        result['errors'] += 1
        except Exception as e:
            result['parse_error'] = str(e)
        return result

    per_file = {Path(p).stem: parse_dumbo_log(p) for p in logs}
    total_rbc_outputs = sum(r.get('rbc_outputs', 0) for r in per_file.values())
    total_aba_decisions = sum(r.get('aba_decisions', 0) for r in per_file.values())
    total_mvba_outputs = sum(r.get('mvba_outputs', 0) for r in per_file.values())
    total_acs_completions = sum(r.get('acs_completions', 0) for r in per_file.values())
    total_consensus_rounds = sum(r.get('consensus_rounds', 0) for r in per_file.values())
    total_errors = sum(r.get('errors', 0) for r in per_file.values())

    # duration from env
    try:
        duration = int(os.environ.get('DURATION_VAL','0'))
        if duration <= 0:
            duration = 1
    except Exception:
        duration = 1

    print('--- DUMBO STATS ---')
    print('per_file_results:', per_file)
    print('total_rbc_outputs:', total_rbc_outputs)
    print('total_aba_decisions:', total_aba_decisions)
    print('total_mvba_outputs:', total_mvba_outputs)
    print('total_acs_completions:', total_acs_completions)
    print('total_consensus_rounds:', total_consensus_rounds)
    print('total_errors:', total_errors)
    print('duration_s:', duration)
    print('rbc_outputs_per_sec:', total_rbc_outputs / duration if duration > 0 else 0)
    print('aba_decisions_per_sec:', total_aba_decisions / duration if duration > 0 else 0)
    print('mvba_outputs_per_sec:', total_mvba_outputs / duration if duration > 0 else 0)
    print('consensus_rounds_per_sec:', total_consensus_rounds / duration if duration > 0 else 0)
    print('---------------')
    
    # Write results to CSV
    import csv
    N = int(os.environ.get('N_VAL','4'))
    result_file = os.path.join(RUN_DIR, 'results/result.csv')
    with open(result_file, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow([
            'timestamp', 'N', 'duration_s', 'total_rbc_outputs', 'total_aba_decisions', 
            'total_mvba_outputs', 'total_acs_completions', 'total_consensus_rounds', 
            'total_errors', 'rbc_outputs_per_sec', 'aba_decisions_per_sec', 
            'mvba_outputs_per_sec', 'consensus_rounds_per_sec', 'per_file_results'
        ])
        writer.writerow([
            os.environ.get('TS',''), N, duration, total_rbc_outputs, total_aba_decisions,
            total_mvba_outputs, total_acs_completions, total_consensus_rounds,
            total_errors, total_rbc_outputs / duration if duration > 0 else 0,
            total_aba_decisions / duration if duration > 0 else 0,
            total_mvba_outputs / duration if duration > 0 else 0,
            total_consensus_rounds / duration if duration > 0 else 0,
            str(per_file)
        ])
    print(f'results written to: {result_file}')
PY

echo "[run_local_dumbo] done."